# -*- coding: utf-8 -*- 

'''
程序说明：东门和东南门各有自己的自编码器网络，开始实现算法，使用单次仿真数据而非平均数据
'''

import numpy as np
from cv2 import cv2
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import time
from datetime import datetime
from StateMapDataset import  FakeDataSet ,FakeAvgDataset, FakeSinglePairDataset
import os,sys
from logger import Logger
from AutoEncoder import BehaviorModelAutoEncoder
import itertools


if __name__ == '__main__':
    TestOrTrain = 'train'

    logfileName = 'log' + str(int(time.time()))+'.txt'
    sys.stdout = Logger(logfileName)
    

    resultDir = '/home/hsc/Research/StateMapPrediction/code/models/EastAndSouth3/resultDir/'#可视化结果保存在这里
    fakeSingleTrainset = FakeSinglePairDataset('/home/hsc/Research/StateMapPrediction/datas/fake/EastGate/data2','/home/hsc/Research/StateMapPrediction/datas/fake/SouthEastGate/data2',train = True)
    fakeSingleTestset = FakeSinglePairDataset('/home/hsc/Research/StateMapPrediction/datas/fake/EastGate/data2','/home/hsc/Research/StateMapPrediction/datas/fake/SouthEastGate/data2',train = False)
    fakeSingleTrainLoader = DataLoader(fakeSingleTrainset,batch_size=4,shuffle=True)
    fakeSingleTestLoader = DataLoader(fakeSingleTestset,batch_size=4,shuffle=True)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print('device = ',device)

    if TestOrTrain =='train':

        EastModel = BehaviorModelAutoEncoder()
        SouthEastModel = BehaviorModelAutoEncoder()

        criterion = nn.MSELoss()
        # optimizer = optim.SGD(model.parameters(),lr = 0.0001,momentum=0.9)
        optimizer = optim.Adam(itertools.chain(EastModel.parameters(),SouthEastModel.parameters()),lr = 0.0001)

        EastModel.to(device)
        SouthEastModel.to(device)

        print('Start training.',time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) )
        start_time = time.time()

        lastTestingLoss = np.inf

        for epoch in range(2000):#500个epoch
            running_loss = running_loss1 = running_loss2 = running_loss3 = 0
            count = 0
            for i,sample in enumerate(fakeSingleTrainLoader):
                count += 1
                E,SE = sample['EStateMap'].to(device), sample['SEStateMap'].to(device)
                optimizer.zero_grad()

                EOut,Ez = EastModel(E)
                SOut,Sz = SouthEastModel(SE)

                loss1 = criterion(EOut,E)
                loss2 = criterion(SOut,SE)
                loss3 = criterion(Ez,Sz)

                loss = loss1 + loss2

                if epoch > 500 and epoch%2 == 0:
                    loss = loss + loss3



                loss.backward()
                optimizer.step()
                
                running_loss += loss.item()
                running_loss1 += loss1.item()
                running_loss2 += loss2.item()
                running_loss3 += loss3.item()

                if count == 10:
                    print('[%d, %5d] training loss: %.3f, E-E recons loss: %.3f, S-S recons loss: %.3f, z-z loss: %.3f' %(epoch + 1, i + 1, running_loss / count,running_loss1/count,running_loss2/count,running_loss3/count))
                    count = 0
                    running_loss = running_loss1 = running_loss2 = running_loss3 = 0
                    
                        
            testing_loss = testing_loss1 = testing_loss2 = testing_loss3 = 0
            count = 0
            for i,sample in enumerate(fakeSingleTestLoader):
                E,SE = sample['EStateMap'].to(device), sample['SEStateMap'].to(device)
                optimizer.zero_grad()

                EOut,Ez = EastModel(E)
                SOut,Sz = SouthEastModel(SE)

                loss1 = criterion(EOut,E)
                loss2 = criterion(SOut,SE)
                loss3 = criterion(Ez,Sz)
                loss = loss1 + loss2 + loss3     

                testing_loss  += loss.item()
                testing_loss1 += loss1.item()
                testing_loss2 += loss2.item()
                testing_loss3 += loss3.item()
                count += 1

                EinSout = SouthEastModel.decoder(Ez)
                SinEout = EastModel.decoder(Sz)

                if i == 0:
                    concatenate = torch.cat([E,SE,EOut,SOut,SinEout,EinSout],0)
                    concatenate = concatenate.detach()
                    concatenate = concatenate.cpu()
                    concatenate = torchvision.utils.make_grid(concatenate,nrow=4,normalize=True,pad_value=255)

                    concatenate = 255 - concatenate.numpy()*255
                    concatenate = np.transpose(concatenate,(1,2,0))
                    imgName = 'Epoch%d.jpg'%epoch
                    imgName = resultDir +imgName
                    cv2.imwrite(imgName,concatenate)
                    pass

            if  epoch > 500 and epoch%2 == 0 and testing_loss < lastTestingLoss:
                lastTestingLoss = testing_loss
                torch.save(EastModel.state_dict(), '/home/hsc/Research/StateMapPrediction/code/models/EastAndSouth3/modelParam/Easemodel.pth')
                torch.save(SouthEastModel.state_dict(), '/home/hsc/Research/StateMapPrediction/code/models/EastAndSouth3/modelParam/SEmodel.pth')

            print('[%d] testing loss: %.3f, E-E recons loss: %.3f, S-S recons loss: %.3f, z-z loss: %.3f' %(epoch + 1, testing_loss / count,testing_loss1/count,testing_loss2/count,testing_loss3/count))
            print('Time is ',time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) ,end = ' ')
            using_time = time.time()-start_time
            hours = int(using_time/3600)
            using_time -= hours*3600
            minutes = int(using_time/60)
            using_time -= minutes*60
            print('running %d h,%d m,%d s'%(hours,minutes,int(using_time)))

    if TestOrTrain == 'test':
        EastModel = BehaviorModelAutoEncoder()
        SouthEastModel = BehaviorModelAutoEncoder()

        EastModel.load_state_dict(torch.load('/home/hsc/Research/StateMapPrediction/code/models/EastAndSouth3/modelParam/Easemodel.pth'))
        SouthEastModel.load_state_dict(torch.load('/home/hsc/Research/StateMapPrediction/code/models/EastAndSouth3/modelParam/SEmodel.pth'))
        EastModel.to(device)
        SouthEastModel.to(device)

        for i,sample in enumerate(fakeSingleTrainLoader):
            E,SE = sample['EStateMap'].to(device), sample['SEStateMap'].to(device)

            Ez = EastModel.encoder(E)
            EinSout = SouthEastModel.decoder(Ez)
            Sz = SouthEastModel.encoder(SE)
            SinEout = EastModel.decoder(Sz)

            concatenate = torch.cat([E,SE,SinEout,EinSout],0)
            concatenate = concatenate.detach()
            concatenate = concatenate.cpu()
            concatenate = torchvision.utils.make_grid(concatenate,nrow=4,normalize=True,pad_value=255)

            concatenate = 255 - concatenate.numpy()*255
            concatenate = np.transpose(concatenate,(1,2,0))
            imgName = '/home/hsc/testing.jpg'
            cv2.imwrite(imgName,concatenate)



            Ez = Ez.detach()
            Ez = Ez.cpu()
            Ez = Ez.numpy()
            Sz = Sz.detach()
            Sz = Sz.cpu()
            Sz = Sz.numpy()

            print(Ez[0,:])
            print(Sz[0,:])

            print('hhh')

            


            

        

        
